# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import os

from absl import app, flags

import pandas
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

FLAGS = flags.FLAGS


def do_plot_sim(df_edge_2, df_edge_5, csv_root, product):
  plt.figure(figsize=(10, 10))
  plt.title('Kraken -> OkexFutures %s unhedge mm simulation result' % product)
  plt.xticks(df_edge_2.index, df_edge_2['ma_window'])
  edge_2, = plt.plot(df_edge_2.index, df_edge_2['pnl/qty'], color='g', marker='o')
  edge_5, = plt.plot(df_edge_2.index, df_edge_5['pnl/qty'], color='b', marker='o')
  plt.xlabel('ma_window')
  plt.ylabel('pnl/qty')
  plt.legend([edge_2, edge_5], ['edge_bp=2', 'dege_bp=5'])
  plt.savefig(os.path.join(csv_root, 'Kraken_%s_sim_result.png' % product))


def plot_sim(df, csv_root, product):
  assert len(df['symbol'].unique()) == 1
  assert df['symbol'][0] == product
  df['pnl/qty'] = df['pnl'] / df['qty']
  df_edge_2 = df.loc[df['edge_bp'] == 2]
  df_edge_5 = df.loc[df['edge_bp'] == 5]
  do_plot_sim(df_edge_2, df_edge_5, csv_root, product)


def read_csv_into_df(csv_root, csv):
  csv_dir = os.path.join(csv_root, csv)
  df = pandas.read_csv(csv_dir, sep=',', header=0)
  return df


def main(argv):
  csv_root = FLAGS.csv_root
  assert csv_root, '--csv_root must be specified.'

  products = ['BCH-USD', 'BTC-USD', 'EOS-USD', 'ETC-USD', 'ETH-USD', 'XRP-USD']
  for product in products:
    sim_df = read_csv_into_df(csv_root, 'Kraken_%s.csv' % product)
    plot_sim(sim_df, csv_root, product)

  return 0


if __name__ == '__main__':
  flags.DEFINE_string('csv_root', None, 'Input csv files root directory.')

  app.run(main)
